package com.sakiko.memory.config;

import com.sakiko.memory.service.ChatAssistant;
import com.sakiko.memory.service.ChatMemoryAssistant;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.TokenCountEstimator;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiTokenCountEstimator;
import dev.langchain4j.service.AiServices;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * projectName: com.sakiko.memory.config
 *
 * @author: sakiko
 * time: 2025/8/28 1:17
 * description: 记忆缓存
 * 出处: https://docs.langchain4j.dev/tutorials/chat-memory
 */
@Configuration
public class LLMConfig {

    @Bean
    public ChatModel chatModel() {
        return OpenAiChatModel.builder()
                .apiKey(System.getenv("ALIBABA_BAILIAN_API_KEY"))
                .modelName("qwen-long")  // 专做长文本
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .build();
    }

    @Bean("chat")
    public ChatAssistant chatAssistant(ChatModel chatModel) {
        return AiServices.create(ChatAssistant.class, chatModel); // 简单创建
    }

    //    https://docs.langchain4j.dev/tutorials/chat-memory#eviction-policy
    @Bean("chat-withMemory-Window") // MessageWindowChatMemory
    public ChatMemoryAssistant chatMemoryAssistant(ChatModel chatModel) {
        return AiServices.builder(ChatMemoryAssistant.class) // 高级创建
                .chatModel(chatModel)
                .chatMemoryProvider(id ->  // 按照 memoryId对应创建了一个 chatMemory
                        MessageWindowChatMemory.withMaxMessages(100) // 限制一百条
                )
                .build();
    }

    @Bean("chat-withMemory-Token") // TokenWindowChatMemory
    public ChatMemoryAssistant chatMemoryAssistant2(ChatModel chatModel) {
        // TokenCountEstimator默认的 token分词器，需要结合 Tokenizer计算 ChatMessage的 token数量
        TokenCountEstimator openAiTokenCountEstimator = new OpenAiTokenCountEstimator("gpt-4");
        return AiServices.builder(ChatMemoryAssistant.class)
                .chatModel(chatModel)
                .chatMemoryProvider(id -> TokenWindowChatMemory.withMaxTokens(1000, openAiTokenCountEstimator))
                .build();
    }

}
